Should we mark this as internal to prevent it from leaking into 4.4 before the feature is fully implemented?

Yes, done, and created https://issues.apache.org/jira/browse/KAFKA-20763 to publish it once the feature is complete

Why do we need to do a special non-blocking call? In the existing logic, if an allocation request is blocked, all existing ProducerBatches become immediately drainable.

Not strictly needed, agreed, it was just to fail fast and avoid what seemed like a "waisted" wait here, exactly because of the fact you mention, that blocking would make all open batches drainable, including this one we want to extend (and this is the trick I was trying to avoid, updated the comment that was misleading).

So with a blocking call, even if we get the pool memory we need in time, the batch would probably be gone/drained, and we would need to start a new batch anyways (so opting for non-blocking just to go straight to the new-batch phase if pool is exhausted when extending).

That being said, there is a case where blocking here would seem better, that's if the pool is exhausted at this point but it frees-up before the batch we're extending is drained. But that one didn't seem common in practice given that the sender drains batches before polling (so by the time memory frees up from a request that compelted, the batch we're waiting to extend would be drained already). Makes sense?

It may take a bit of time for the closed batches to be drained. If we continue here, it seems that the client will just busy-loop until some batches are drained and some free space becomes available in buffer pool?

even without the drain the closeForAppends should be enough to make that the next iteration find that the stream is closed => isFull => !hasRoom so starting a new batch (blocking, no busy loop). Makes sense? Added a test testExhaustedExtensionFallsBackToBlockingNewBatchPath to show the behaviour.

It's a bit awkward to have a return value of null and RecordAppendResult.needsExtension. Could we introduce a non-null value to indicate the batch is full?

Agreed, done. Introduced needsNewBatch (aligned with needsBufferExtension)

I guess two concurrent clients could add buffers exceeding the batch size? Those buffers won't be used, but can only be freed after the batch is drained.

Correct, there could be over-reservation at this point with concurrent appends. As for how the extra buffers would be used/released, I see the following cases:

• used right away: next loop iteration -> the loser's tryAppend right after this line fails with needsExtension (the winner already took the free capacity that was used to calculate this extension). Then the loser loops and tries again (recomputes the extension needed, so it uses what it had added concurrently + more)
• left unused, freed when batch drains (case you mentioned) -> if the loser's record can't land in this batch (batch closes and loser's record goes into new batch).

Added 2 tests to cover this: testConcurrentExtensionRaceDoesNotOverflowChunkedStream and testBatchSizeLimitRespectedDespiteOverReservation

Here is a related case. Two concurrent producer each allocated its own extensionChunks. It's possible that the first producer's allocation is enough to fit the records from both producer. It would be useful to document that this logic optimizes for the common case, but can lead to slightly temporary over allocation.

Yes, agreed, added a comment describing it.

This is a no-op since all pooled chunks have been used if we reach here.

agreed, removed.

Why is remainingBytes guaranteed to be an int? memoryRequired could be larger than int and pooled.size() initially could be 0.

yup, good catch, seemed safe here but it's not because of how memoryRequired is rounded up higher up. This code here does not exist anymore though so no changes (the other place where a similar unsafe op was being done to calculate stillNeeded went away too, removed with the comment/fix below)

This may be ok, but it's a bit weird to free up the chunks only to be reallocated again. Here is an alternative that doesn't require a freeup() call.

  // Reuse pooled chunks first. If a reused chunk covers a slot we already reserved as                                                                                                                                                                                                                                    
  // raw bytes in an earlier iteration, hand that raw reservation back to the pool.                                                                                                                                                                                                                                       
  while (pooled.size() < numChunks && !free.isEmpty()) {                                                                                                                                                                                                                                                                  
      pooled.add(free.pollFirst());                                                                                                                                                                                                                                                                                       
      if (accumulated >= chunkSize) {          // accumulated is always chunk-aligned here                                                                                                                                                                                                                                
          accumulated -= chunkSize;                                                                                                                                                                                                                                                                                       
          this.nonPooledAvailableMemory += chunkSize;   // refund → available to other waiters                                                                                                                                                                                                                            
      }                                                                                                                                                                                                                                                                                                                   
  }                                                                                                                                                                                                                                                                                                                       
  // Reserve raw memory for any still-uncovered chunks, in whole chunks.                                                                                                                                                                                                                                                  
  while (pooled.size() + (int)(accumulated / chunkSize) < numChunks                                                                                                                                                                                                                                                       
          && this.nonPooledAvailableMemory >= chunkSize) {                                                                                                                                                                                                                                                                
      this.nonPooledAvailableMemory -= chunkSize;                                                                                                                                                                                                                                                                         
      accumulated += chunkSize;                                                                                                                                                                                                                                                                                           
  }                                                                                                                                                                                                                                                                                                                       

nice, makes sense, done.

The first condition is redundant, given the second one.

 (pooled+1)*chunkSize + accumulated ≤ memoryRequired                                                                                                                                                                                                                                                               
    ⇒  (pooled+1)*chunkSize ≤ memoryRequired − accumulated ≤ memoryRequired = numChunks*chunkSize                                                                                                                                                                                                                         
    ⇒  pooled+1 ≤ numChunks                                                                                                                                                                                                                                                                                               
    ⇒  pooled < numChunks

Agreed (the condition changed though, with the comment above, and does not have the redundancy anymore, so no changes directly for this)

Could we return accumulated and pooled chunks in the same place? For example, we can set a flag like allocationCompleted to replace accumulated = 0. Then we can free both accumulated and pooled chunks if the flag is not set.

Done. Introduced the allocationCompleted flag and now all refunds happen together in a finally, much better indeed.

Would it better to use FULL and INCREMENTAL?

pending => pendingNewBatch?

size => firstAppendSize?

Why does it need to sleep for 50ms, which is long for a unit test? Ditto below.

yeap not needed in this case really, and the other one can be better done with a waitForCondition on the available memory, done.

slow-success => chunked-waiter ?

If we switch automatically to full because of the batch size, it might be useful to log a warning.

yes, done

Hmm, this still seems to only work if remainingBytesRequired is 1. If it's larger than 1 and the current chunk doesn't have enough remaining bytes, we move to the next chunk. A subsequent write will write on the new chunk, the remaining bytes on the previous chunk are wasted.

yeap, agreed, fixed. Also added a todo to review this bit when compression and the need to grow lands

At this point it's expected that the allocated chunks have capacity

Hmm, is this true? This method could return null because the batch is full, right?

right, updated the javadoc to clarify

Does this still need to be protected? Ditto for other fields. We can also remove the addition of BufferPool to spotbugs-exclude.xml.

agreed, all private now and cleaned up the spotbugs

This can be private.

This can have package level visibility.

This can be private.

This can be private.

Should we validate that recordsBuilder.bufferStream() is chunked?

yes, done

It seems that we can block on the bufferpool longer than maxTimeToBlock. Claude suggested the following path that doesn't exist in RecordAccumulator.

Iteration A: needsNewBatch → blocking allocateChunks(_, maxTimeToBlock) succeeds after a wait. A concurrent thread created a batch meanwhile, so appendNewBatch's in-lock tryAppend returns needsBufferExtension → line 257 deallocates the stream, continue.
Iteration B: extension path does the non-blocking allocateChunks(gap, 0L) → BufferExhaustedException → closes the batch (line 168) → continue.
Iteration C: batch now closed → needsNewBatch → blocks again for a fresh full maxTimeToBlock.

fixed, tracking remainingTimeToBlock along the whole path (which goes as you described above)

Iteration A consumes from it, B does not touch it because it's non-blocking, C consumes from what's left (fixing the gap)

int numChunks = (totalSize - 1) / chunkSize + 1; is overflow-safe even without the cast. The existing code is correct, so feel free to ignore this comment.

remainingTimeToBlock = Math.max(0L, remainingTimeToBlock - Math.max(0L, nowMs - allocationStartMs)); ?

This adds an extra time.milliseconds() call on the critical path. Could you do some perf test between full and incremental to verify there is no degradataion?

Long story short, the extra time doesn't seem to have a relevant impact on append, but the repeated record size calculation on extensionBytesNeeded does, related to what we discussed in the comment above, but the numbers I ran now do show the impact even without headers. I included a partial improvement here now, and will add more as a follow-up with https://issues.apache.org/jira/browse/KAFKA-20859 (just avoiding it in this PR as it moves quite a few things)

Now the longer story, I ran a JMH comparison of the two accumulators append (will add the benchmarks in a separate PR)

1 - Batch creation path (every append to a different partition, so every one creates a batch): the extra time.milliseconds accounts for only 3.5% of the difference between the strategies at default batch.size (and it's per batch, not per record, above code is the new batch path). Overall on this path incremental is +21% at batch.size=16KB but −42% at bigger batch size 256KB (cold pool, not "steady" state)

2 - Steady-state append (records landing in an existing batch) is where the numbers show an impact, with 100-byte records: +36% at default batch.size and +45% at bigger batch.size 256KB. At 16KB, extensionBytesNeeded is around 91% of that gap (the record gets sized twice per append). The partial fix I added in this PR is to scoped the capacity check in ChunkedProducerBatch.tryAppend to a batch's first record only (the only append where the caller does not validate so we need to). That improves it to +17.6% at 16KB and +37.6% at 256KB. I will close the gap with what we know for now (follow-up with KAFKA-20859 right away), and run it all again there to see where we get.

Makes sense?